#coding:utf-8


import argparse
import os
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage

def mkdir(path):
    full_path=os.path.join(os.getcwd(),path)
    #print(full_path)
    if not os.path.exists(full_path):
        os.mkdir(full_path)

def reshapes(kernel,num_rows_image,num_cols_image):
    '''
           改变Kernel的shape,
           将kernel中的参数分不到矩阵 kernel——reshape=len(kernel)*(num_cols_image*num_rows_image)
           最后计算卷积使用矩阵乘法
    :param kernel: 卷积核 np.array
    :param num_rows_image: 输入图片的行
    :param num_cols_image: 输入图片的高
    :return:
    '''
    print('reshapes')
    rows_outputs=1+(num_rows_image-kernel.shape[0])
    cols_outputs=1+(num_cols_image-kernel.shape[1])
    num_len=num_cols_image*num_rows_image
    kernel_reshape=np.zeros(shape=(rows_outputs*cols_outputs,num_len))
    for row_offest in range(rows_outputs*cols_outputs):
        for row in range(kernel.shape[0]):
            for col in range(kernel.shape[1]):
                kernel_reshape[
                    row_offest,
                    row*rows_outputs+row_offest+col
                ]=kernel[row,col]
    return kernel_reshape,rows_outputs,cols_outputs


def perform_conv(image,kernel):
    '''
    卷积操作，将卷积操作转化为两个大矩阵乘法
    :param image:
    :param kernel:
    :return:
    '''
    print('conving')
    kernel_reshape,rows_output,cols_output=reshapes(
        kernel,
        image.shape[0],
        image.shape[1])
    image_reshape=image.reshape((image.shape[0]*image.shape[1],1))

    # print('Kernel\t:',kernel)
    # print('image shape\t:',np.shape(image))#（200,200）
    # print('kernel_reshape shape\t:',kernel_reshape.shape)#39204,40000
    # print('rows',rows_output)#198
    # print('cols',cols_output)#198
    return np.matmul(
        kernel_reshape,image_reshape).reshape(rows_output,cols_output)

def perform_trans_conv(matrix,kernel,rows_output_matrix,cols_output_matrix):
   '''
    卷积操作（反卷积）
   :param matrix: 输入的图像
   :param kerneal: 卷积卷积核
   :param rows_output_matrix: 输出图像的row
   :param cols_output_matrix: 输出图像的col
   :return:
   '''
   print('transpose conving')
   kernel_reshape,_,_=reshapes(kernel,rows_output_matrix,cols_output_matrix)
   matrix_reshape=matrix.reshape((matrix.shape[0]*matrix.shape[1],1))
   return np.matmul(
       np.transpose(kernel_reshape),
       matrix_reshape
   ).reshape(rows_output_matrix,cols_output_matrix)

def play(image_path,num_convs):
    i=0
    #创建输出目录
    folder='examples'
    mkdir(folder)
    output_folder=folder+'/{}_{:02d}_convs'.format(
        image_path.split(os.path.sep)[-1],
        num_convs
    )
    mkdir(output_folder)
    image=ndimage.imread(image_path).mean(axis=2)
    # print(image.shape)
    # print(image.mean(axis=2).shape)
    plt.pcolor(image)
    # plt.show()
    print('saving origin image')
    plt.savefig(os.path.join(output_folder,'{:02d}_origin_image'.format(i)))

    #定义kernels
    kernels=[]
    for _ in range(num_convs):
        kernels.append(np.random.randint(low=0,high=3,size=9).reshape((3,3)))

    #执行卷积操作
    convlotions=[image]
    for j in  range(num_convs):
        i+=1
        convlotions.append(perform_conv(convlotions[-1],kernels[j]))
        plt.pcolor(convlotions[-1])
        plt.savefig(os.path.join(output_folder, '{:02d}_convs_{:02}.png'.format(i,j+1)))

    print('conv ended ')
    #反卷积操作
    transpose_convlutions=[convlotions[-1]]
    for j in range(num_convs):
        i+=1
        transpose_convlutions.append(perform_trans_conv(
            transpose_convlutions[-1],
            kernels[-1-j],
            convlotions[-2-j].shape[0],
            convlotions[-2-j].shape[1]
        ))
        plt.pcolor(transpose_convlutions[-1])
        plt.savefig(os.path.join(output_folder, '{:02d}_convs_trans_{:02}.png'.format(i, j + 1)))

def parse_args():
    parser=argparse.ArgumentParser(description='test convlution and transpose convlution')
    parser.add_argument('--image',dest='image',type=str,default='../data/python2.png')
    parser.add_argument('--num_convs',dest='num_convs',type=int,default=3)
    return parser.parse_args()

def main():
    args=parse_args()
    play(args.image,args.num_convs)

if __name__=="__main__":
    main()



